package com.shxs.LLMConnect.BlueLM;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.shxs.LLMConnect.LLMConnector;
import com.zhipu.oapi.Constants;
import com.zhipu.oapi.service.v4.model.ChatCompletionRequest;
import com.zhipu.oapi.service.v4.model.ChatMessage;
import com.zhipu.oapi.service.v4.model.ChatMessageRole;
import com.zhipu.oapi.service.v4.model.ModelApiResponse;
import org.springframework.http.*;
import org.springframework.web.client.RestTemplate;

import java.util.*;

public class BlueLM implements LLMConnector {
    String appId = "3033296844";
    String appKey = "vVMiQybXRSxxhqOo";
    String URI = "/vivogpt/completions";
    String DOMAIN = "api-ai.vivo.com.cn";
    String METHOD = "POST";

    String messages = "";
    public String ask(String messages) {

        try {
            UUID requestId = UUID.randomUUID();

            Map<String, Object> map = new HashMap<>();
            map.put("requestId", requestId.toString());
            String queryStr = mapToQueryString(map);

            //构建请求体
            Map<String, Object> data = new HashMap<>();

            data.put("messages", JSONArray.parse(messages));
            data.put("model", "vivo-BlueLM-TB");
            UUID sessionId = UUID.randomUUID();
            data.put("sessionId", sessionId.toString());
            System.out.println(sessionId);


            HttpHeaders headers = VivoAuth.generateAuthHeaders(appId, appKey, METHOD, URI, queryStr);
            headers.add("Content-Type", "application/json");
            System.out.println(headers);
            String url = String.format("http://%s%s?%s", DOMAIN, URI, queryStr);
            String requsetBodyString = new ObjectMapper().writeValueAsString(data);
            RestTemplate restTemplate = new RestTemplate();
            HttpHeaders httpHeaders = new HttpHeaders();
            httpHeaders.setContentType(MediaType.valueOf(MediaType.APPLICATION_JSON_VALUE));
            httpHeaders.addAll(headers);
            HttpEntity<String> requestEntity = new HttpEntity<>(requsetBodyString, httpHeaders);

            ResponseEntity<String> response = restTemplate.exchange(url, HttpMethod.POST, requestEntity, String.class);
            if (response.getStatusCode() == HttpStatus.OK) {
                System.out.println("Response body: " + response.getBody());
                JSONObject js = JSON.parseObject(response.getBody());
                String line = js.getString("code");
                if (!line.equals("0")) return null;
                line = js.getJSONObject("data")
                        .getString("content");
                return line;
            } else {
                System.out.println("Error response: " + response.getStatusCode());
            }
            return null;
        }
        catch (Exception e) {
            System.out.println(e.toString());
            return null;
        }
    }

    public static String mapToQueryString(Map<String, Object> map) {
        if (map.isEmpty()) {
            return "";
        }
        StringBuilder queryStringBuilder = new StringBuilder();
        for (Map.Entry<String, Object> entry : map.entrySet()) {
            if (queryStringBuilder.length() > 0) {
                queryStringBuilder.append("&");
            }
            queryStringBuilder.append(entry.getKey());
            queryStringBuilder.append("=");
            queryStringBuilder.append(entry.getValue());
        }
        return queryStringBuilder.toString();
    }

    public int getHistoryLength(){
        return messages.length();
    };
    public String askWithHistory(String line){
        String chatMessage = "{\n" +
                "    \"role\": \"user\",\n" +
                "    \"content\": \""+line+"\"\n" +
                "  },";
        messages+=chatMessage;
        line=ask("["+messages.substring(0, messages.length()-1)+"]");
        chatMessage="{\n" +
                "    \"role\": \"assistant\",\n" +
                "    \"content\": \""+line+"\"\n" +
                "  },";
        messages+=chatMessage;
        return line;
        //return "["+messages.substring(0, messages.length()-1)+"]";
    };
}
